# -*- coding: utf-8 -*-
"""
Core implementation of the proposed model:
- SBERT-enhanced TextRank for keyword extraction
- Source Credibility Scoring (SCS)
- Multi-level Transformer (BERT + Longformer)
- Multi-task Learning (Bias Classification + Ranking)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import BertModel, LongformerModel
from sentence_transformers import SentenceTransformer
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

# ----------------------------
# 1. SBERT-enhanced TextRank
# ----------------------------
class SBERTTextRank:
    def __init__(self, model_name='all-MiniLM-L6-v2'):
        self.model = SentenceTransformer(model_name)

    def extract_keywords(self, sentences, top_k=10):
        vectors = self.model.encode(sentences)
        sim_matrix = cosine_similarity(vectors)

        scores = np.ones(len(sentences))
        damping = 0.85
        for _ in range(20):
            for i in range(len(sentences)):
                scores[i] = (1 - damping) + damping * sum(
                    sim_matrix[i][j] / sum(sim_matrix[j]) * scores[j] for j in range(len(sentences)) if i != j
                )

        ranked_indices = np.argsort(scores)[::-1][:top_k]
        return [sentences[i] for i in ranked_indices]

# ----------------------------
# 2. Source Credibility Score
# ----------------------------
def compute_scs(domain_score, platform_score, history_score, alpha=0.4, beta=0.3, gamma=0.3):
    scs = alpha * domain_score + beta * platform_score + gamma * history_score
    return np.tanh(scs)

# ----------------------------
# 3. Multi-level Transformer
# ----------------------------
class MultiLevelTransformer(nn.Module):
    def __init__(self, bert_name='bert-base-uncased', longformer_name='allenai/longformer-base-4096'):
        super(MultiLevelTransformer, self).__init__()
        self.bert = BertModel.from_pretrained(bert_name)
        self.longformer = LongformerModel.from_pretrained(longformer_name)

        hidden_size = self.bert.config.hidden_size
        self.gate = nn.Linear(hidden_size * 2, hidden_size)
        self.classifier = nn.Linear(hidden_size, 2)  # Binary classification
        self.regressor = nn.Linear(hidden_size, 1)   # Ranking score

    def forward(self, input_ids_bert, attention_mask_bert, input_ids_long, attention_mask_long):
        out_bert = self.bert(input_ids=input_ids_bert, attention_mask=attention_mask_bert).last_hidden_state[:, 0, :]
        out_long = self.longformer(input_ids=input_ids_long, attention_mask=attention_mask_long).last_hidden_state[:, 0, :]

        gate_weight = torch.sigmoid(self.gate(torch.cat([out_bert, out_long], dim=1)))
        fused = gate_weight * out_bert + (1 - gate_weight) * out_long

        bias_logits = self.classifier(fused)
        ranking_score = self.regressor(fused)

        return bias_logits, ranking_score

# ----------------------------
# 4. Loss Functions
# ----------------------------
class MultiTaskLoss(nn.Module):
    def __init__(self, lambda_weight=0.5):
        super(MultiTaskLoss, self).__init__()
        self.lambda_weight = lambda_weight
        self.ce = nn.CrossEntropyLoss()
        self.rank_loss = nn.MarginRankingLoss(margin=1.0)

    def forward(self, bias_logits, bias_labels, rank_scores_pos, rank_scores_neg):
        loss_class = self.ce(bias_logits, bias_labels)
        loss_rank = self.rank_loss(rank_scores_pos, rank_scores_neg, torch.ones_like(rank_scores_pos))
        return self.lambda_weight * loss_class + (1 - self.lambda_weight) * loss_rank